from datetime import datetime,timedelta,date
from chinese_calendar import is_workday


class Trade_date():
    def __init__(self,base_day=None):
        if not base_day:
            self.base_day = datetime.today().date()
        self.set_base_day(base_day)
    
    def get_date(self,days,end_day=None):
        """计算获取days个交易日的开始日期

        Args:
            end_day: 截止时间,输入为空时默认为今天
            days: 交易日天数
        """
        # 格式化时间参数
        if not end_day:
            end_day = datetime.today().date() # 今天
        elif type(end_day) == str and end_day.isdigit() and len(end_day) == 8:
            end_day = datetime.strptime(end_day, '%Y%m%d').date()
        else:
            print('输入日期格式错误')
            return
        # 设置基础时间
        self.set_base_day(end_day)
        # 获取迭代器
        last_day = self.last_trade_day()
        # 迭代days-1次
        for i in range(days-1):
            next(last_day)
        
        return next(last_day),end_day.strftime('%Y%m%d')
    
    def set_base_day(self,base_day):
        """设置基准日期
        """ 
        if type(base_day) == date:
            self.base_day = base_day
        elif type(base_day) == str and base_day.isdigit() and len(base_day) == 8:
            self.base_day = datetime.strptime(base_day, '%Y%m%d').date()
        else:
            return
        
    def next_trade_day(self):
        """迭代获取基准日的下一个交易日
            yield 返回迭代器
        """
        one_day = timedelta(days=1) # 正1天
        today =self.base_day
        # 如果当天是交易日，先返回当天
        if today.isoweekday() < 6 and is_workday(today):
            yield today.strftime('%Y%m%d')
        # 循环迭代
        while(True):
            today = today + one_day
            if today.isoweekday() < 6 and is_workday(today):
                # 如果不是周六日，且不是节假期（包括调休）,那就是交易日
                rs_date = today.strftime('%Y%m%d')
                # 如果这一天在未来，返回None
                if rs_date > datetime.now().strftime('%Y%m%d'):
                    yield None
                else:
                    yield rs_date
    
    def last_trade_day(self):
        """迭代获取基准日的上一个交易日
            yield 返回迭代器
        """
        one_day = timedelta(days=-1) # 负1天
        today =self.base_day
        # 如果当天是交易日，先返回当天
        if today.isoweekday() < 6 and is_workday(today) and datetime.now().hour > 17:
            yield today.strftime('%Y%m%d')
        # 循环迭代
        while(True):
            today = today + one_day
            if today.isoweekday() < 6 and is_workday(today):
                # 如果不是周六日，且不是节假期（包括调休）,那就是交易日
                yield today.strftime('%Y%m%d')
                
    def is_trade_day(self,date):
        assert type(date) == str and date.isdigit() and len(date) == 8,'格式不正确,要求YYYYMMdd'
        date = datetime.strptime(date, '%Y%m%d').date()
        if date.isoweekday() < 6 and is_workday(date):
            return True
        return False